package com.lwc.cfdns.aspect;

import com.lwc.cfdns.annotation.RateLimit;
import com.lwc.cfdns.constants.Constants;
import com.lwc.cfdns.enums.LimitType;
import com.lwc.cfdns.utils.RedisUtils;
import com.lwc.cfdns.utils.ip.IpUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.annotation.Order;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.Collections;

/**
 * 限流切面
 *
 * @author luwc
 * @title RateLimit
 * @description
 * @time 2023/8/10 15:00:43
 * @version: 1.0
 */
@Aspect
@Order(10) // 定义切面执行顺序
@Component
public class RateLimitAspect {
    private final Logger logger = LoggerFactory.getLogger(RateLimitAspect.class);

    @Resource
    private RedisUtils redisUtils;

    /**
     * 定义切面
     * 所有Controller
     * 或手动注解
     *
     * @param
     * @return void
     * @throws
     * @version 1.0.0
     * @author luwc
     * @time 2023/8/10 15:21:40
     */
    @Pointcut("execution(public * com.lwc.cfdns.controller.*.*.*.*(..)) || @annotation(com.lwc.cfdns.annotation.RateLimit)")
    public void rateLimit() {

    }

    /**
     * 环绕通知,环绕增强，相当于MethodInterceptor
     * 需要注意:@annotation(xxx) 时才可以跟第二个参数
     *
     * @param proceedingJoinPoint
     * @return Object
     * @throws
     * @version 1.0.0
     * @author luwc
     * @time 2023/8/10 16:23:51
     */
    @Around("rateLimit()")
    public Object aroundRateLimit(ProceedingJoinPoint proceedingJoinPoint) throws Throwable {
        Object rtValue = null;
        try {
            String whitelistKey = Constants.RATE_LIMIT_WHITELIST_KEY;
            String lockKey = Constants.RATE_LIMIT_LOCK_KEY;
            String ipAddr = IpUtils.getIpAddr();
            if (redisUtils.hasKey(whitelistKey) && redisUtils.sIsMember(whitelistKey, ipAddr)) {
                // 白名单内
                // 重要，调用目标函数执行 反回为结果
                rtValue = proceedingJoinPoint.proceed();
                return rtValue;
            }
            if (redisUtils.hasKey(lockKey + ipAddr)) {
                // 已锁定的IP
                throw new RuntimeException("请求太频繁，请稍后再试。");
            }
            // 获取注解中的属性
            MethodSignature signature = (MethodSignature) proceedingJoinPoint.getSignature();
            RateLimit rateLimit = signature.getMethod().getAnnotation(RateLimit.class);
            String key = Constants.RATE_LIMIT_KEY;
            long limit = 100l;
            long expire = 60l;
            long lockExpire = 3600l;
            if (rateLimit != null) {
                key = rateLimit.key();
                limit = rateLimit.limit();
                expire = rateLimit.expire();
                lockExpire = rateLimit.lockExpire();
            }

            StringBuffer stringBuffer = new StringBuffer(key);
            // 类名
            String className = proceedingJoinPoint.getTarget().getClass().getName();
            // 方法名
            String methodName = proceedingJoinPoint.getSignature().getName();

            stringBuffer.append(className).append("-").append(methodName);
            if (rateLimit != null && rateLimit.limitType() == LimitType.IP) {
                stringBuffer.append("-").append(ipAddr);
            }
            key = stringBuffer.toString();
            Long number = redisUtils.getRedisTemplate().execute(limitScript(), Collections.singletonList(key), limit, expire);
            if (number != null && number.longValue() != 0l && number.longValue() <= limit) {
                logger.info("限流时间段内访问了第：{} 次,限制次数：{},缓存key：{}", number.toString(), limit, key);
                // 重要，调用目标函数执行 反回为结果
                rtValue = proceedingJoinPoint.proceed();
                return rtValue;
            }
            // 锁IP处理
            redisUtils.set(lockKey + ipAddr, 1, lockExpire);
            throw new RuntimeException("请求太频繁，请稍后再试。");
        } catch (Throwable e) {
            e.getStackTrace();
            throw new RuntimeException(e.getMessage());
        } finally {
        }
    }

    /**
     * 自定义lua脚本
     *
     * @param
     * @return String
     * @throws
     * @version 1.0.0
     * @author luwc
     * @time 2023/8/10 15:12:17
     */
    private RedisScript<Long> limitScript() {
        String stringScript = "local key = KEYS[1]\n" +
                "local limit = tonumber(ARGV[1])\n" +
                "local expireTime = tonumber(ARGV[2])\n" +
                "local current = tonumber(redis.call('get', key) or '0')\n" +
                "if current + 1 > limit then\n" +
                "   return 0\n" +
                "else\n" +
                "   redis.call('INCRBY', key, '1')\n" +
                "   redis.call('expire', key, expireTime)\n" +
                "   return current + 1\n" +
                "end";
        //指定ReturnType为Long.class，只支持List.class, Boolean.class和Long.class
        RedisScript<Long> redisScript = new DefaultRedisScript<>(stringScript, Long.class);
        return redisScript;
    }
}
